/*
 * Copyright (c) 2024 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef NET_FIREWALL_CONNTRACK_H
#define NET_FIREWALL_CONNTRACK_H

#include <linux/compiler.h>
#include "netfirewall_utils.h"
#include "netfirewall_conntrack_def.h"
#include "netfirewall_conntrack_map.h"

static __always_inline void ct_reset_seen_flags(struct ct_entry *entry)
{
    entry->rx_flags_seen = 0;
    entry->tx_flags_seen = 0;
}

static __always_inline void ct_reset_closing(struct ct_entry *entry)
{
    entry->rx_closing = 0;
    entry->tx_closing = 0;
}

static __always_inline bool ct_entry_alive(const struct ct_entry *entry)
{
    return !entry->rx_closing || !entry->tx_closing;
}

static __always_inline bool ct_entry_closing(const struct ct_entry *entry)
{
    return entry->tx_closing || entry->rx_closing;
}

static __always_inline enum ct_action ct_tcp_select_action(union tcp_flags flags)
{
    if (flags.value & (TCP_FLAG_RST | TCP_FLAG_FIN))
        return ACTION_CLOSE;

    if (flags.value & TCP_FLAG_SYN)
        return ACTION_CREATE;

    return ACTION_UNSPEC;
}

static __always_inline bool ct_entry_seen_both_syns(const struct ct_entry *entry)
{
    bool rx_syn = entry->rx_flags_seen & TCP_FLAG_SYN;
    bool tx_syn = entry->tx_flags_seen & TCP_FLAG_SYN;

    return rx_syn && tx_syn;
}

static __always_inline bool __ct_update_timeout(struct ct_entry *entry, __u32 lifetime, enum ct_dir dir,
    union tcp_flags flags, __u8 report_mask)
{
    __u32 now = bpf_mono_now();
    __u8 accumulated_flags;
    __u8 seen_flags = flags.lower_bits & report_mask;
    __u32 last_report;

    WRITE_ONCE(entry->lifetime, now + lifetime);

    if (dir == CT_INGRESS) {
        accumulated_flags = READ_ONCE(entry->rx_flags_seen);
        last_report = READ_ONCE(entry->last_rx_report);
    } else {
        accumulated_flags = READ_ONCE(entry->tx_flags_seen);
        last_report = READ_ONCE(entry->last_tx_report);
    }
    seen_flags |= accumulated_flags;
    if (last_report + bpf_sec_to_mono(CT_REPORT_INTERVAL) < now || accumulated_flags != seen_flags) {
        /* verifier workaround: we don't use reference here. */
        if (dir == CT_INGRESS) {
            WRITE_ONCE(entry->rx_flags_seen, seen_flags);
            WRITE_ONCE(entry->last_rx_report, now);
        } else {
            WRITE_ONCE(entry->tx_flags_seen, seen_flags);
            WRITE_ONCE(entry->last_tx_report, now);
        }
        return true;
    }
    return false;
}
/**
 * @brief Update the CT timeouts for the specified entry.
 *
 * @param entry struct ct_entry
 * @param dir enum ct_dir
 * @param seen_flags union tcp_flags
 * @return If CT_REPORT_INTERVAL has elapsed since the last update, updates the
 *         last_updated timestamp and returns true. Otherwise returns false.
 */
static __always_inline bool ct_update_timeout(struct ct_entry *entry, bool tcp, enum ct_dir dir,
    union tcp_flags seen_flags)
{
    __u32 lifetime = bpf_sec_to_mono(CT_CONNECTION_LIFETIME_NONTCP);
    bool syn = seen_flags.value & TCP_FLAG_SYN;

    if (tcp) {
        entry->seen_non_syn |= !syn;
        if (entry->seen_non_syn) {
            lifetime = bpf_sec_to_mono(CT_CONNECTION_LIFETIME_TCP);
        } else {
            lifetime = bpf_sec_to_mono(CT_SYN_TIMEOUT);
        }
    }

    return __ct_update_timeout(entry, lifetime, dir, seen_flags, CT_REPORT_FLAGS);
}

/**
 * @brief create a key pair of ct_tuple/ct_entry and add to conntrack map
 *
 * @param tuple struct ct_tuple
 * @param skb struct __sk_buff
 * @param dir enum ct_dir
 * @return true if success or false if an error occurred
 */
static __always_inline bool ct_create(struct ct_tuple *tuple, struct __sk_buff *skb, const enum ct_dir dir)
{
    /* Create entry in original direction */
    struct ct_entry entry = {};
    bool is_tcp = tuple->protocol == IPPROTO_TCP;
    union tcp_flags seen_flags = {
        .value = 0
    };

    seen_flags.value |= is_tcp ? TCP_FLAG_SYN : 0;
    ct_update_timeout(&entry, is_tcp, dir, seen_flags);

    return bpf_map_update_elem(&CONNTRACK_MAP, tuple, &entry, 0) == 0;
}

/**
 * @brief lookup from conntrack map by ct_tuple if found then update lifetime of connection
 *
 * @param skb struct __sk_buff
 * @param tuple struct ct_tuple
 * @param dir enum ct_dir
 * @return CT_NEW if not found, otherwise CT_RELATED, CT_REOPENED or CT_ESTABLISHED
 */
static __always_inline enum ct_status ct_lookup(struct __sk_buff *skb, const struct ct_tuple *tuple, enum ct_dir dir)
{
    struct ct_entry *entry = bpf_map_lookup_elem(&CONNTRACK_MAP, tuple);
    if (entry) {
        __u32 l3_nhoff = get_l3_nhoff(skb);
        bool is_tcp = is_l4_protocol(skb, l3_nhoff, IPPROTO_TCP);
        union tcp_flags seen_flags = {};
        if (is_tcp) {
            __u32 l4_nhoff = get_l4_nhoff(skb);
            if (load_tcp_flags(skb, l4_nhoff, &seen_flags) < 0) {
                return CT_RELATED;
            }
        }
        if (ct_entry_alive(entry)) {
            ct_update_timeout(entry, is_tcp, dir, seen_flags);
        }
        enum ct_action action = ct_tcp_select_action(seen_flags);
        switch (action) {
            case ACTION_CREATE:
                if (ct_entry_closing(entry)) {
                    ct_reset_closing(entry);
                    ct_reset_seen_flags(entry);
                    entry->seen_non_syn = false;
                    ct_update_timeout(entry, is_tcp, dir, seen_flags);
                    return CT_REOPENED;
                }
                break;
            case ACTION_CLOSE:
                /* If we got an RST and have not seen both SYNs,
                 * terminate the connection. (For CT_SERVICE, we do not
                 * see both directions, so flags of established
                 * connections would not include both SYNs.)
                 */
                if (!ct_entry_seen_both_syns(entry) && (seen_flags.value & TCP_FLAG_RST)) {
                    entry->rx_closing = 1;
                    entry->tx_closing = 1;
                } else if (dir == CT_INGRESS) {
                    entry->rx_closing = 1;
                } else {
                    entry->tx_closing = 1;
                }

                if (ct_entry_alive(entry)) {
                    break;
                }
                __ct_update_timeout(entry, bpf_sec_to_mono(CT_CLOSE_TIMEOUT), dir, seen_flags, CT_REPORT_FLAGS);
                break;
            default:
                break;
        }
        return CT_ESTABLISHED;
    }

ct_new:
    return CT_NEW;
}

#endif // NET_FIREWALL_CONNTRACK_H